package com.hlz.flink.chapter06;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;

import java.util.concurrent.TimeUnit;

/**
 * 自定义水位线产生
 *
 * @author Hongliang Zhu
 * @create 2023-06-18 22:26
 */
public class CustomWatermarkTest {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<String> dataStreamSource = env.socketTextStream("127.0.0.1", 9696);
        SingleOutputStreamOperator<Stock> stockStream = dataStreamSource.filter(StringUtils::isNotEmpty)
                .map(new MapFunction<String, Stock>() {
                    @Override
                    public Stock map(String value) throws Exception {
                        final String[] split = value.split(",");
                        return new Stock(split[0], Double.parseDouble(split[1]), split[2], Long.parseLong(split[3]));
                    }
                });


        SingleOutputStreamOperator<Stock> stockSingleOutputStreamOperator = stockStream.assignTimestampsAndWatermarks(new MyCustomWatermarkStrategy());

        stockSingleOutputStreamOperator.keyBy(value -> value.getId())
                .window(TumblingEventTimeWindows.of(Time.of(10, TimeUnit.MILLISECONDS)))
                .maxBy("price")
                .print("custom_watermark");

        env.execute();


        /*

         stock_001,1,UDFStock,1
         stock_002,8,UDFStock,8
         stock_003,10,UDFStock,10
         stock_003,12,UDFStock,12
         stock_007,7,UDFStock,7
         stock_004,22,UDFStock,22
         stock_003,8,UDFStock,10
         stock_003,8,UDFStock,10


         */
    }


    /**
     * 自定义水位线策略
     */
    public static class MyCustomWatermarkStrategy implements WatermarkStrategy<Stock> {

        @Override
        public WatermarkGenerator<Stock> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
            return new CustomPeriodicGenerator();
        }

        @Override
        public TimestampAssigner<Stock> createTimestampAssigner(TimestampAssignerSupplier.Context context) {
            return new SerializableTimestampAssigner<Stock>() {
                @Override
                public long extractTimestamp(Stock element, long recordTimestamp) {
                    return element.getTimeStamp();
                }
            };
        }
    }


    /**
     * 自定义水位线生成器
     */
    public static class CustomPeriodicGenerator implements WatermarkGenerator<Stock> {
        private Long delayTime = 5000L; // 延迟时间

        // 观察到的最大时间戳
        private Long maxTs = Long.MIN_VALUE + delayTime + 1L;

        @Override
        public void onEvent(Stock event, long eventTimestamp, WatermarkOutput output) {
            // 没来一条数据就调用一次，更新一下最大的时间戳
            maxTs = Math.max(event.getTimeStamp(), maxTs);
            System.out.println("update maxTs: "+maxTs+" -->event time: " +  eventTimestamp);

        }

        @Override
        public void onPeriodicEmit(WatermarkOutput output) {
            // 发射水位线，默认 200ms 调用一次
            output.emitWatermark(new Watermark(maxTs - delayTime - 1L));
            System.out.println("onPeriodicEmit--->" +(maxTs - delayTime - 1L));
        }
    }

}
